package cn.com.idmy.auth.token.jwt;

import cn.com.idmy.auth.login.LoginState;
import cn.com.idmy.auth.token.TokenManager;
import cn.com.idmy.base.exception.UnauthorizedException;
import cn.com.idmy.base.util.Assert;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.auth0.jwt.JWT;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.exceptions.SignatureVerificationException;
import com.auth0.jwt.exceptions.TokenExpiredException;
import com.auth0.jwt.interfaces.DecodedJWT;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;
import lombok.NonNull;
import lombok.experimental.Accessors;
import org.dromara.hutool.core.convert.ConvertUtil;
import org.dromara.hutool.core.lang.tuple.Pair;
import org.dromara.hutool.core.text.StrUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.util.Objects;

@Data
@NoArgsConstructor
@Accessors(fluent = true)
@EqualsAndHashCode(callSuper = true)
public class JwtTokenManager<ID> extends TokenManager<ID> {
    public static String ID = "id";
    public static final String AT = "at";
    public static final String TC = "tc";
    public static String NAME = "name";
    public static final String DOT = ".";

    public synchronized void init() {
        super.init();
        tokenCreator = this::createToken;
    }

    @NotNull
    public String createToken(@NotNull Algorithm algorithm, @NotNull Duration exp, long cfgId, long accountType, @NotNull Object id, @Nullable JSONObject payload) {
        var now = Instant.now();
        return cfgId + DOT + JWT.create()
                .withPayload(payload)
                .withClaim(ID, (Long) id)
                .withClaim(AT, accountType)
                .withClaim(TC, cfgId)
                .withExpiresAt(now.plus(exp))
                .withIssuedAt(now)
                .sign(algorithm);
    }

    @NotNull
    public String createToken(@NotNull Algorithm algorithm, @NotNull Duration exp, @NotNull String json) {
        return JWT.create().withPayload(json).withExpiresAt(Instant.now().plus(exp)).sign(algorithm);
    }

    public Pair<Long, String> parseToken(String token) {
        var arr = token.split("\\.");
        var cfgStr = arr[0];
        token = token.substring(cfgStr.length() + 1);
        var cfgId = Assert.notNull(ConvertUtil.toLong(cfgStr, null), "令牌配置不能为空");
        return new Pair<>(cfgId, token);
    }

    @NotNull
    public LoginState<ID> parseToken(@NotNull Algorithm algorithm, long cfgId, @NotNull String token) {
        DecodedJWT jwt;
        try {
            jwt = JWT.require(algorithm).build().verify(token);
        } catch (TokenExpiredException e) {
            throw new UnauthorizedException("令牌已过期");
        } catch (SignatureVerificationException e) {
            throw new UnauthorizedException("令牌签名错误");
        } catch (Exception e) {
            throw new UnauthorizedException("令牌错误");
        }

        var cfg = jwt.getClaim(TC).asLong();
        if (cfg != null && !Objects.equals(cfg, cfgId)) {
            throw new UnauthorizedException("令牌配置ID非法");
        }

        @SuppressWarnings("unchecked")
        var id = (ID) jwt.getClaim(ID).asLong();

        var state = new LoginState<>(id);
        var at = jwt.getClaim(AT).asLong();
        if (at != null) {
            state.accountType(at);
        }
        var name = jwt.getClaim(NAME).asString();
        if (StrUtil.isNotBlank(name)) {
            state.data().put(NAME, name);
        }

        var zone = ZoneId.systemDefault();
        var expires = jwt.getExpiresAtAsInstant();
        var issued = jwt.getIssuedAtAsInstant();

        state.expAt(expires.atZone(zone).toLocalDateTime());
        if (issued != null) {
            state.createdAt(issued.atZone(zone).toLocalDateTime());
        }
        return state;
    }

    @Override
    public @NotNull LoginState<ID> getState(@NotNull String token) {
        return parseAccessToken(token);
    }

    protected JSONObject getAccessTokenPayloadData(LoginState<ID> state) {
        return JSONObject.of(NAME, state.data().getString(NAME));
    }

    protected JSONObject getRefreshTokenPayloadData(LoginState<ID> state) {
        return null;
    }

    protected void createToken(@NotNull LoginState<ID> state) {
        var cfg = state.config();
        var tokenAlgorithm = (Algorithm) cfg.tokenAlgorithm();

        var token = createToken(tokenAlgorithm, cfg.tokenExp(), cfg.id(), state.accountType(), state.id(), getAccessTokenPayloadData(state));
        state.token(token);

        var refreshTokenAlgorithm = (Algorithm) cfg.refreshTokenAlgorithm();
        if (StrUtil.isBlank(state.refreshToken()) && refreshTokenAlgorithm != null) {
            var exp = Assert.notNull(cfg.refreshTokenExp(), "刷新令牌有效期不能为空");
            var refreshToken = createToken(refreshTokenAlgorithm, exp, cfg.id(), state.accountType(), state.id(), getRefreshTokenPayloadData(state));
            state.refreshToken(refreshToken);
        }

        var idTokenAlgorithm = (Algorithm) cfg.idTokenAlgorithm();
        if (StrUtil.isBlank(state.idToken()) && idTokenAlgorithm != null) {
            state.idToken(createToken(idTokenAlgorithm, cfg.tokenExp(), JSON.toJSONString(state.data())));
        }
    }

    protected @NotNull LoginState<ID> parseAccessToken(@NonNull String token) {
        var pair = parseToken(token);
        var cfgId = pair.getLeft();
        var cfg = tokenConfigDao.getTokenConfig(cfgId);
        var algorithm = (Algorithm) cfg.tokenAlgorithm();
        LoginState<ID> state = parseToken(algorithm, cfgId, pair.getRight());
        state.config(cfg);
        state.token(token);
        return state;
    }
}
